Skip to content

fix(grpo): clip the KL log-ratio before exp to avoid inf overflow - #6637

Open
behroozazarkhalili wants to merge 11 commits into
mainfrom
fix/3015-kl-exp-overflow-v2
Open

fix(grpo): clip the KL log-ratio before exp to avoid inf overflow#6637
behroozazarkhalili wants to merge 11 commits into
mainfrom
fix/3015-kl-exp-overflow-v2

Conversation

@behroozazarkhalili

@behroozazarkhalili behroozazarkhalili commented Aug 3, 2026

Copy link
Copy Markdown
Collaborator

What

The K3 KL estimator overflows to inf when the policy and reference distributions drift far apart during training (a large positive log-ratio), as reported in #3015:

per_token_kl = torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1

Fix

Add an opt-in GRPOConfig.kl_log_ratio_clip (float, optional, default None) that clips the log-ratio before the exponential:

kl_log_ratio = ref_per_token_logps - per_token_logps
if self.args.kl_log_ratio_clip is not None:
    clip = torch.tensor(self.args.kl_log_ratio_clip, dtype=kl_log_ratio.dtype, device=kl_log_ratio.device)
    if not torch.isfinite(torch.exp(clip)):
        raise ValueError(...)  # the clip itself must keep exp finite in the working dtype
    kl_log_ratio = kl_log_ratio + (kl_log_ratio.clamp(max=clip) - kl_log_ratio).detach()
per_token_kl = torch.exp(kl_log_ratio) - kl_log_ratio - 1

The clip is upper-only (a large negative log-ratio underflows exp to zero and leaves K3 finite) and straight-through: the value is clamped but the gradient passes as if it were not. A plain clamp zeroes the K3 slope for clipped tokens, and with use_bias_correction_kl=True (the default) the term reduces to K3(clip) * ratio, whose gradient rewards a lower policy log-prob, so minimizing the loss would push the policy further from the reference. Measured with plain torch at x = 20, clip = 10: d(KL)/d(log-prob) is +2.2e4 with a plain clamp, -10 straight-through, and -32 unclipped in float32, where the exact -x * ratio = -20 is already lost to cancellation.

  • When None (the default) the estimator is bit-identical to the current one, so there is no behavior change for existing users.
  • When set, it keeps per_token_kl finite while preserving the K3 shape (exp(x) - x - 1 stays non-negative for the clamped x). The value must be positive and finite: GRPOConfig rejects the rest, since a non-positive clip invents KL at an exact match and -inf reaches the estimator as inf past the trainer's overflow guard.

This is approach (a) from @albertvillanova's guidance on #3015 (clip the log-ratio before exp behind an opt-in field).

Consistency

The K3 block is duplicated in two experimental trainers that subclass GRPOTrainer (gmpo, gspo_token); the same guard is applied to both so the trainers stay aligned per the repo's duplication policy.

Tests

Seven kl_log_ratio_clip tests in tests/test_grpo_trainer.py, each building a tiny GRPOTrainer on CPU and calling _compute_loss directly on hand-built inputs (so they exercise the real K3 branch, not a standalone helper): the unclipped estimator overflows to inf above the dtype's exp ceiling and clipping tames it; a clip above that ceiling raises; a large negative log-ratio is left intact; a normally scaled log-ratio is unaffected; the gradient of the loss with respect to a clipped token's log-prob stays negative with the bias correction on and off; the KL term with zero advantages stays finite and positive; and five non-positive or non-finite clips are rejected. A copy of the tree with a plain clamp fails both gradient cases, and one with the validation removed fails all five rejection cases. ruff check + ruff format clean.

Since this is a numerical-stability guard rather than a new paper method, no paper_index.md entry was added.

Resolves #3015

cc @qgallouedec @kashif @albertvillanova


Note

Medium Risk
Touches the KL penalty in core GRPO loss paths; default-off behavior limits blast radius, but enabled clipping changes gradients for extreme policy–reference drift.

Overview
Fixes #3015 by adding optional GRPOConfig.kl_log_ratio_clip: when set, the K3 KL term upper-bounds ref − policy log-ratio with a straight-through clamp before exp, so large positive ratios no longer blow up to inf. Default None leaves the estimator unchanged.

GRPOConfig validates the clip (positive, finite) and rejects values whose own exp would overflow the working dtype; GRPOTrainer raises NotImplementedError if use_liger_kernel=True with non-zero beta and a clip set, since the fused Liger path cannot apply it.

The same K3 clipping block is mirrored in gmpo and gspo_token experimental trainers. tests/test_grpo_trainer.py adds regression coverage (overflow vs clip, validation, one-sided behavior, gradients with bias correction, Liger incompatibility).

Reviewed by Cursor Bugbot for commit 612c558. Bugbot is set up for automated code reviews on this repo. Configure here.

The K3 KL estimator exp(ref - cur) - (ref - cur) - 1 overflows to inf when the policy and reference distributions drift far apart during training (large positive log-ratio).

Add an opt-in GRPOConfig.kl_log_ratio_clip (float, optional, default None) that clips the log-ratio to [-clip, clip] before the exponential. When None the estimator is bit-identical to the current one, so there is no behavior change for existing users. The same guard is applied to the two experimental trainers that duplicate the K3 block (gmpo, gspo_token) to keep them consistent, and a regression test covers the overflow and the no-op default.

Resolves #3015
Comment thread trl/trainer/grpo_trainer.py Outdated
@bot-ci-comment

bot-ci-comment Bot commented Aug 3, 2026

Copy link
Copy Markdown

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@qgallouedec

Copy link
Copy Markdown
Member

Thanks, I'm not strongly opposed, but afaik, no-one is using a ref model regularization anymore. So it'm not sure if adding a new parameter again is really worth it

The kl_log_ratio_clip guard reused the name log_ratio inside the beta != 0 block, shadowing the importance-sampling log-ratio (per_token_logps - old_per_token_logps) that is used later in _compute_loss (log_ratio_per_token) and in GMPO's clip-fraction metrics. Rename the KL estimator's local to kl_log_ratio so the IS log-ratio is preserved.

@cursor cursor Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

Fix All in Cursor

❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.

Reviewed by Cursor Bugbot for commit 707808f. Configure here.

Comment thread trl/trainer/grpo_trainer.py
@behroozazarkhalili

Copy link
Copy Markdown
Collaborator Author

Fair point, and I would rather not add public config surface you have to keep stable if the KL path is winding down.

The crash itself is still real for anyone running beta != 0 (the K3 exp(ref - cur) overflows to inf on entropy collapse), but it can be fixed without a new parameter. A few lines up, the sequence-importance path already clamps a log-ratio unconditionally:

# grpo_trainer.py, existing
log_ratio_clamped = torch.clamp(log_ratio_per_token, -20.0, 20.0)

So I can drop kl_log_ratio_clip and instead clamp the K3 log-ratio to a fixed [-20, 20], consistent with the above. exp(20) is about 4.8e8, so it is a no-op for any legitimate per-token KL and only bites in the overflow regime. That fixes the crash, adds no new public API, and keeps the three duplicated K3 blocks (grpo, gmpo, gspo_token) consistent.

Want me to switch #6637 to that (unconditional clamp, no new field)? If you would rather not touch the KL path at all given its declining use, I am equally happy to just close this. cc @albertvillanova

@behroozazarkhalili

Copy link
Copy Markdown
Collaborator Author

Follow-up: Bugbot just flagged that kl_log_ratio_clip is silently ignored on the Liger path. The fused LigerFusedLinearGRPOLoss computes the K3 KL inside the kernel, so a Python-side clip never runs, and there is no init-time guard for the combination.

That is actually another reason to prefer the unconditional-clamp route: with no public kl_log_ratio_clip field there is no option that can silently no-op on Liger. The non-Liger path still gets the overflow fix, and the Liger kernel's internal K3 stays a separate upstream concern. So dropping the parameter resolves both your config-surface point and the Liger finding in one go. Happy to push that version whenever you give the nod.

…silently ignoring it

The Liger fused GRPO loss computes the KL penalty internally and does not receive
`kl_log_ratio_clip`, so enabling the clip under `use_liger_kernel=True` silently left the
issue #3015 `inf`-overflow guard inactive. Raise `NotImplementedError` at init (matching the
existing Liger-incompatible-option guards) when `beta != 0` and `kl_log_ratio_clip` is set,
and add a `require_liger_kernel` test asserting the raise.
The `kl_log_ratio_clip` added for issue #3015 clamped the K3 estimator
on both sides and promised to keep the KL term finite for any positive
value. Neither held.

Clamping the negative side changes a correct number. `exp` overflows
only for large positive input; a large negative log-ratio underflows
to zero and leaves K3 finite, growing as `-kl_log_ratio - 1`. At
`kl_log_ratio = -50` with `kl_log_ratio_clip = 10` the estimator
dropped from 49.0 to 9.0, with no error and no warning. The clamp is
now `max=` only.

The finiteness promise failed above the working dtype's exp
ceiling. float32 overflows at 88.7229, so `kl_log_ratio_clip=90.0`
clamped to 90 and still returned `inf`. That ceiling varies by
dtype: roughly 88.7 for float32 and bfloat16, but only 11.1 for
float16. Deriving it as `math.log(finfo.max)` is itself unsafe,
because the float64 result rounds up when cast back down and the
bound then overflows anyway. The clip is instead exponentiated as a
scalar in the tensor's own dtype, and a value that overflows raises
`ValueError` rather than silently returning `inf`.

The regression test never called the trainer. It reimplemented K3
in a local helper and asserted against that helper, so deleting the
clamp from all three trainers left it at 2 passed, exit 0. Four cases
now drive `_compute_loss` with hand-built inputs, one per property:
overflow tamed by a usable clip, an over-large clip rejected,
a large negative log-ratio left intact, and a non-binding clip as
a no-op. The dtype ceiling each case needs is read from the tensor
rather than written as a literal, so the same test is correct under
bfloat16 autocast and float16.

Three mutants confirm the tests bite: deleting the guard fails 2 of
the 4, restoring the two-sided clamp fails 2, and removing only the
`ValueError` fails 1. The clamp block stays byte-identical across grpo,
gmpo and gspo_token.
@behroozazarkhalili

Copy link
Copy Markdown
Collaborator Author

Verification on GPU, since the Liger path cannot be exercised on a CPU runner.

Ran on one H100 at the pushed head 390356b9, torch 2.11.0 with triton 3.6.0, in two passes.

The 13 Liger nodes: 13 passed, exit 0, in 232s. These are the tests that go red locally without a GPU, where triton raises "0 active drivers" before any assertion runs, so a red there says nothing about the code. On real hardware they pass.

The four kl_log_ratio_clip cases (5 selected): 5 passed, exit 0, in 53s. That covers the clip staying finite above the dtype's exp threshold and the one-sided clamp leaving a large negative log-ratio alone.

This is the same hardware question raised on the Bugbot thread about the Liger path. The fused loss still raises NotImplementedError when kl_log_ratio_clip is set, and that guard is among the 13.

…ts value

A plain clamp on the log-ratio zeroes the K3 slope for every clipped token. With the bias correction
on (the default) the KL term then reduces to `K3(clip) * ratio`, and the ratio's gradient rewards a
lower policy log-prob: minimizing the loss pushes the policy further from the reference. Measured with
plain torch at `x = 20`, `clip = 10`: the gradient with respect to the policy log-prob is `+2.2e4`
under the clamp, `-10` with the clip straight-through, and `-32` unclipped in float32, where the true
`-x * ratio = -20` is already lost to cancellation. The clip now keeps its value but passes the
gradient, in all three copies of the block, so a clipped token still pulls the policy toward the
reference with the slope at the clip.

`GRPOConfig` accepted any float: a non-positive clip invents KL at an exact policy/reference match, and
`-inf` reaches the estimator as `inf` past the trainer's overflow guard, which only rejects values whose
exponential overflows. `__post_init__` now requires a positive finite number. The docstring and help
said a policy drifting far above the reference overflows; a large positive `log(pi_ref / pi_theta)`
means the policy sits far below it.

Tests: the gradient of the loss with respect to a clipped token's log-prob must be negative, with the
bias correction on and off; the KL term with zero advantages must stay finite and positive; five
non-positive or non-finite clips must be rejected.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Numerical stability in KL's exp function

2 participants